import pickle
from collections.abc import Iterator
from typing import TYPE_CHECKING, TypeVar

from zkl_aiutils_training import Resumable, RunningTaskPlugin

from zkl_ptutils_training.plugins.resumable_fs import FsPauseArgs, FsResumeArgs, FsResumeFromCheckpointArgs
from .processing import MLProducer

if TYPE_CHECKING:
    from zkl_aiutils_datasets import Dataset

AnyInput = TypeVar('AnyInput')


class ZklDatasetProducer(MLProducer[AnyInput], RunningTaskPlugin, Resumable):
    def __init__(self, dataset: 'Dataset[AnyInput]'):
        self._dataset = dataset
        self._iterator: Iterator | None = None

    def on_resume(self, args: FsResumeArgs):
        if isinstance(args, FsResumeFromCheckpointArgs):
            from zkl_aiutils_datasets import resume
            with args.checkpoint_fs.open('iterator.pkl', 'rb') as fp:
                state = pickle.load(fp)
            self._iterator = resume(self._dataset, state)

    def on_pause(self, args: FsPauseArgs):
        if args.checkpoint_fs is not None:
            from zkl_aiutils_datasets import pause
            state = pause(self._iterator)
            with args.checkpoint_fs.open('iterator.pkl', 'wb') as fp:
                pickle.dump(state, fp)

    # producer

    def on_before_run(self):
        if self._iterator is None:
            self._iterator = iter(self._dataset)

    def produce(self) -> AnyInput:
        return next(self._iterator)
